from __future__ import division
from numpy import histogram2d, histogram
import Image
import ImageFilter
import ImageMath
import Preprocessing as Pre
import numpy as np
import pylab
import scipy.ndimage as ndi
import time
"""
Image metric module
"""

epsilon = 2.2e-16 


def GetValue(img1, img2, method='mse', gsigma=1.4, roi=None, ret_histo=False, bins=256, debug=False):
    """
    cost = cost_function(x, optfunc_args)    --- OR ---
    cost, joint_histogram = cost_function(x, optfunc_args)   

    computes the alignment between 2 2D images using one of several 
    available metrics. Before the metric is computed, the images are histogram
    equalized.  The 8 bit scaling is done using an integrated histogram method and
    is called prior to this.
    
    The parameters to be supplied will vary depending on the metric to be used, 
    they are not all required for all metrics, see details below.
    
    method='gd',roi=None,bins=256,gsigma=1.4,ret_histo=False,debug=False
    
    Parameters 
    ----------

        method [All Metrics]  : {'nmi', 'mi', 'ncc', 'ecc', 'mse', 'ncc' 'gd', 'gc', 'bpc'}
            flag for type of registration metric. 
            Statistical Group:
                nmi  - normalized mutual information; 
                mi   - mutual information; 
                ecc  - entropy cross correlation;
                necc - normalized entropy cross correlation
            Other (no histograms) 
                ncc  - normalized cross correlation. 
                mse  - mean square error. 
                gd   - gradient difference (Penney,)
                gc   - gradient correlation
                bpc  - MY METHOD (gc+mi)
        gsigma ['mse', 'gd] : {float} 
            Gaussian sigma
        roi [All Metrics] : 2x2 Array specifying the top left corner and the 
            bottom right corner of the region of the image that the metric
            should be applied to.  [[TopRightX,TopRightY],[BottomX,BottomY]]
            If no value is given, the whole image is used.
            The Top left image coordinate is 0,0 and the bottom right 
            coordinate is image.shape 
        
        ret_histo ['nmi', 'mi', 'ncc', 'ecc'] : {0, 1} 
            if 0 return is: cost 
            if 0 return is: cost, joint_histogram  

    Returns 
    -------
        cost : {float}
            the negative of one of the mutual information metrics
            or negative cross correlation. use negative as the optimization
            is minimization.

        --- OR --- (if ret_histo = 1)

        cost : {float}
            the negative of one of the mutual information metrics
            or negative cross correlation. use negative as the optimization
            is minimization.

        joint_histogram : {nd_array}
            the 2D (256x256) joint histogram of the two volumes
    
    """
#===============================================================================
#    General image processing, common for all metrics
#===============================================================================
    cost = 0.0
    
    img1_original = img1.copy()
    img2_original = img2.copy()
    
    if roi == None:
        # If no region of interest is set, use the whole image.
        # roi is a bounding box, with the first element being the top 
        # left corner and the second element being the bottom right corner
        roi = [[0, 0], img1.shape]
    
    # Convert roi to an integer array
    roi = np.asarray(roi,'int')
    
    # Convert from RGB to grayscale (if necessary)
    # Returns a numpy array
    img1 = Pre.make_gray(img1)
    img2 = Pre.make_gray(img2)
    
    # Crop the images to the roi (both images use same roi
    img1 = img1[roi[0,0]:roi[1,0], roi[0,1]:roi[1,1]]
    img2 = img2[roi[0,0]:roi[1,0], roi[0,1]:roi[1,1]]
    
#===============================================================================
#   Begin metric selection statements 
#===============================================================================
    if img1.shape != img2.shape:
        print "Cannot compute, images must be the same size for this metric!"
        return 0
    
    if method == 'mse':
        # mean squared error method
        cost = FindMSE(img1, img2, gsigma=gsigma, debug=debug)    
        return cost
    elif method == 'bpc':
        # my custom metric
        cost1 = FindGC(img1,img2,gsigma,debug)
        #print 'GC: ', cost1
        
        #8bit images are required for further processing steps (histeq,pdf,cdf)
        #This function has been moved to the specific metrics that need it since it
        #may adversely affect some.
        img1 = Pre.ConvertTo8BitImage(img1)
        img2 = Pre.ConvertTo8BitImage(img2)
        
        # Set a range, so that intensity values outside will be ignored
        # this is useful to eliminate extreme values (eg. 0 and 255)
        r = np.array([[1,bins-2],[1,bins-2]])
        # Do a type of statistical histogram metric        
        joint_histogram, xedges, yedges = \
            np.histogram2d(img2.flatten(), img1.flatten(), (bins, bins), range = r, normed = True)
        
        joint_histogram = joint_histogram + epsilon # prevent log(0) 
        # normalize the joint histogram
        #joint_histogram /= joint_histogram.sum() 
        # get the marginals
        marginal_col = joint_histogram.sum(axis=0)
        marginal_row = joint_histogram.sum(axis=1)
        
        cost2 = FindMI(joint_histogram,marginal_col,marginal_row)
        #print 'MI: ', cost2
        return (0.5*cost1+0.5*cost2)
    
    elif method == 'ssd':
        # Sum squared distance
        #cost = 
        return cost
    elif method == 'ncc':
        # Normalized Cross Correlation
        cost = FindNCC(img1,img2)
        return cost
    elif method == 'nc':
        # Normalized Correlation Metric from ITK
        cost = FindNC(img1,img2)
        return cost
    elif method == 'gd':
        # Calculate gradient difference
        cost = FindGD(img1, img2, gsigma=gsigma, debug=debug)
        return cost
    elif method == 'gc':
        # Calculate gradient correlation metric
        cost = FindGC(img1,img2,gsigma,debug)
        return cost
    elif (method == np.asarray(['nmi', 'mi', 'necc', 'ecc'])).any():
        #8bit images are required for further processing steps (histeq,pdf,cdf)
        #This function has been moved to the specific metrics that need it since it
        #may adversely affect some.
        img1 = Pre.ConvertTo8BitImage(img1)
        img2 = Pre.ConvertTo8BitImage(img2)
        
        # Set a range, so that intensity values outside will be ignored
        # this is useful to eliminate extreme values (eg. 0 and 255)
        r = np.array([[1,bins-2],[1,bins-2]])
        # Do a type of statistical histogram metric        
        joint_histogram, xedges, yedges = \
            np.histogram2d(img2.flatten(), img1.flatten(), (bins, bins), range = r, normed = True)
        
        joint_histogram = joint_histogram + epsilon # prevent log(0) 
        # normalize the joint histogram
        #joint_histogram /= joint_histogram.sum() 
        # get the marginals
        marginal_col = joint_histogram.sum(axis=0)
        marginal_row = joint_histogram.sum(axis=1)
        
        if method == 'mi':
            # mutual information
            cost = FindMI(joint_histogram,marginal_col,marginal_row)   
        elif method == 'ecc':
            # entropy correlation coefficient 
            cost = FindECC(joint_histogram,marginal_col,marginal_row)
        elif method == 'nmi':
            # normalized mutual information
            cost =  FindNMI(joint_histogram,marginal_col,marginal_row)
        elif method == 'necc':
            # cross correlation from the joint histogram 
            cost = FindNECC(joint_histogram,marginal_col,marginal_row)
        
        if debug:
            DebugStats(joint_histogram,marginal_col,marginal_row,bins)      
        if ret_histo:
            return cost, joint_histogram 
        else:
            return cost
        
        
#================================================================================
#    Begin Metric Algorithms Code
#================================================================================   
        
def FindGD(img1, img2, gsigma=1.4, scale=None, debug=False):
    """
    Compute Gradient Difference metric.
    
    img1,img2   - numpy image arrays 
    scale       - scalar value to multiply to im2 (this will be determined by 
                  the metric automatically if None is passed
    return      - value (scalar)
    """
    # When gradient difference is first called, the scale between images must be calculated
    if scale == None:
        #Apply a gaussian blur to the image
        img1 = ndi.gaussian_filter(img1, gsigma)
        img2 = ndi.gaussian_filter(img2, gsigma)
        
        #Find the image gradient for the fixed image, using the sobel filter
        img1 = ndi.generic_gradient_magnitude (img1, ndi.sobel)
        img2 = ndi.generic_gradient_magnitude (img2, ndi.sobel)
    
        #Find the image differences, s is a scale factor
        #print "scale: ", s
        scale = FindScale(img1, img2, FindGD)
    
    # Find difference image between two gradient images    
    Idiff = img1 - scale * img2
    A = np.var(Idiff)
    cost = np.sum(A / (A + Idiff ** 2))
    
    if debug:
        print "Scale :", scale
        print "Edge1 (max/min): ", img1.max(), "/", img1.min()
        print "Edge2 (max/min): ", img2.max(), "/", img2.min()
        pylab.figure()
        pylab.subplot(311)
        pylab.imshow(img1, hold=False)
        pylab.title('Image1 edge')
        pylab.subplot(312)
        pylab.imshow(img2, hold=False)
        pylab.title('Image2 edge')
        pylab.subplot(313)
        pylab.imshow(np.abs(Idiff), hold=False)
        pylab.title('Diff Image')
    
    return cost

def FindScale(img1, img2, fx):
    """
    Find the optimal scale value so that images can be subtracted 
    leaving a minimum pixel value.  this is essentially a 1D optimization 
    problem
    
    img1, img2 - numpy image arrays
    fx       - function to optimize (minimize).  Must have call signature:
               fx(img1,img2,scale)
    return s - scale value.  When multiplied with im2, produces minimum value of fx 
    """
    #The starting minimum value is s = 0
    best = 0
    minVal = 10e32
    r = 100
    
    #The step size is the ratio of the intensity ranges of the two images
    #divided by 100
    stepSize = (img1.max() / img2.max()) / r
    for s in range(1, r, 1):
        #loop until a minimum is found
        scale = s * stepSize
        cost = fx(img1, img2, scale=scale, debug=False)
        if cost < minVal:
            minVal = cost
            best = scale
    return best

def FindGC(img1, img2, gsigma, debug=False):
    #Apply a gaussian blur to the image
    img1 = ndi.gaussian_filter(img1, gsigma)
    img2 = ndi.gaussian_filter(img2, gsigma)
    
    #Find the image gradient for the fixed image, using the sobel filter
    img1 = ndi.generic_gradient_magnitude (img1, ndi.sobel)
    img2 = ndi.generic_gradient_magnitude (img2, ndi.sobel)
    
    cost = FindNCC(img1,img2,debug)
    return cost

def FindMSE(img1, img2, gsigma=1.4, debug=False):
    """
    Compute mean square error metric
    """
    #blur images
    img1 = Pre.blur(img1, gsigma)
    img2 = Pre.blur(img2, gsigma)
    
    cost = (np.square(img1 - img2)).mean()
    if debug:
        fig = pylab.figure()
        fig.add_subplot(221)
        pylab.imshow(img1)
        pylab.title('Image 1')
        fig.add_subplot(222)
        pylab.imshow(img2)
        pylab.title('Image 2')
        fig.add_subplot(212)
        mse_img = np.square(img1 - img2)
        pylab.imshow(mse_img)
        pylab.title('MSE Image')
        print 'blur: ', gsigma
        print 'Max MSE pix: ', mse_img.max()
        print 'Metric Value: ', cost
    
    # cost is min when img1 and img2 are aligned so keep cost positive
    return cost

def FindNCC(img1,img2,debug=False):
    """
    Compute Normalized Cross Correlation.
    
    input:
    img1, img2    -     numpy arrays
    
    output:
    cost          -     scalar (float)
    
    From Hipwell. IEEE Transaction on Medical Imaging, Vol 22, No 11, 2003, pg 1425.
    When images are perfectly aligned, will be equal to 1
    """

    img1 = img1.astype('f')
    img2 = img2.astype('f')
    
    cost = (np.sum((img1 - img1.mean())*(img2 - img2.mean())) / 
           (np.sqrt(np.sum((img1 - img1.mean())**2)) * 
            np.sqrt(np.sum((img2 - img2.mean())**2))))
    if debug:
        fig = pylab.figure()
        fig.add_subplot(221)
        pylab.imshow(img1)
        pylab.title('Image 1')
        fig.add_subplot(222)
        pylab.imshow(img2)
        pylab.title('Image 2')
        fig.add_subplot(212)
        ncc_img = ((img1 - img1.mean())*(img2 - img2.mean()) / 
                  (np.sqrt((img1 - img1.mean())**2) * 
                   np.sqrt((img2 - img2.mean())**2)))
        pylab.imshow(ncc_img)
        pylab.title('NCC Image')
        print 'Max NCC pix: ', ncc_img.max()
        print 'Metric Value: ', cost
    
    return - cost

def FindNC(img1,img2):
    """
    Normalized Correlation Metric
    The itk::NormalizedCorrelationImageToImageMetric computes pixel-wise crosscorrelation
    and normalizes it by the square root of the autocorrelation of the images:
    NC(A,B) = -1 * SUM (Ai * Bi)/(SQRT(SUM(Ai**2)) * SUM(Bi**2))
    Ai is the i-th pixel of Image A
    Bi is the i-th pixel of Image B
    N is the number of pixels considered
    Note the -1 factor in the metric computation. This factor is used to make the metric be optimal
    when its minimum is reached. The optimal value of the metric is then minus one. Misalignment
    between the images results in small measure values. The use of this metric is limited to images
    obtained using the same imaging modality. The metric is insensitive to multiplicative factors
    between the two images. This metric produces a cost function with sharp peaks and well defined
    minima. On the other hand, it has a relatively small capture radius.
    """
    img1 = img1.astype('f')
    img2 = img2.astype('f')
    
    cost = -1 * ((np.sum(img1*img2 )) / 
                (np.sqrt(np.sum(img1**2) * np.sum(img2**2))))
    
    return cost
    

def FindMI(joint_histogram,marginal_col,marginal_row):
    """
    Calculate Mutual Information metric
    Described in many references, eg Eq 3.25 (pg 60) Medical Image Registration
    """    
    marginal_outer = np.outer(marginal_row, marginal_col)
    H = joint_histogram * np.log(joint_histogram / marginal_outer)  
    mutual_information = H.sum()
    cost = - mutual_information
    
    return cost

def FindECC(joint_histogram,marginal_col,marginal_row):
    """
    Calculate Entropy Cross Correlation metric
    """ 
    mutual_information = -FindMI(joint_histogram,marginal_col,marginal_row)
    row_entropy = marginal_row * np.log(marginal_row)
    col_entropy = marginal_col * np.log(marginal_col)
    ecc = - 2.0 * mutual_information / (row_entropy.sum() + col_entropy.sum())
    cost = - ecc
    return cost

def FindNMI(joint_histogram,marginal_col,marginal_row):
    """
    Calculate Normalized Mutual Information
    Using Eq 3.30 (pg 61) from Medical Image Registration
    """
    row_entropy = marginal_row * np.log(marginal_row)
    col_entropy = marginal_col * np.log(marginal_col)
    H = joint_histogram * np.log(joint_histogram)  
    nmi = (row_entropy.sum() + col_entropy.sum()) / (H.sum())
    cost = - nmi
    return cost

def FindNECC(joint_histogram,marginal_col,marginal_row):
    """
    Calculate Normalized Entropy Cross Correlation
    """
    r, c = joint_histogram.shape
    i = np.array(range(1, c + 1))
    j = np.array(range(1, r + 1))
    m1 = (marginal_row * i).sum()
    m2 = (marginal_col * j).sum()
    sig1 = np.sqrt((marginal_row * (np.square(i - m1))).sum())
    sig2 = np.sqrt((marginal_col * (np.square(j - m2))).sum())
    [a, b] = np.mgrid[1:c + 1, 1:r + 1]
    a = a - m1
    b = b - m2
    # element multiplies in the joint histogram and grids
    H = ((joint_histogram * a) * b).sum()
    ncc = H / (np.dot(sig1, sig2)) 
    cost = - ncc
    return cost


def DebugStats(joint_histogram,marginal_col,marginal_row,bins):
    """
    Useful information output to debug statistical metrics 
    """
    # Find Product of marginal pdf's
    marginal_outer = np.outer(marginal_row, marginal_col)
    #marginal_outer = marginal_outer + epsilon
    
    # Find Entropy
    HA = -np.sum(marginal_row * np.log(marginal_row))
    HB = -np.sum(marginal_col * np.log(marginal_col))
    
    # Find Joint Entropy
    HAB = -np.sum(joint_histogram * np.log(joint_histogram))
    
    # Find Conditional Entrpoy
    HAbarB = HAB - HA
    HBbarA = HAB - HB
    
    # Find Mutual Information
    IAB1 = joint_histogram * np.log(joint_histogram/marginal_outer)
    IAB2 = HA + HB - HAB
    IAB3 = HA - HBbarA  
    IAB4 = HB - HAbarB
    
    print "Mutual Information diagnostics:"
    print "    I(A,B)1 = ", IAB1.sum()
    print "    I(A,B)2 = ", IAB2
    print "    I(A,B)3 = ", IAB3
    print "    I(A,B)4 = ", IAB4
    
    pylab.figure()
    pylab.subplot(221)
    pylab.imshow(joint_histogram, origin='lower')
    pylab.title('Joint Histogram')
    pylab.subplot(222)
    pylab.plot(marginal_col)
    pylab.plot(marginal_row)
    pylab.title('Marginal Histograms')
    pylab.subplot(223)
    pylab.imshow(joint_histogram * np.log(joint_histogram), origin='lower')
    pylab.title('Joint Entropy')
    pylab.subplot(224)
    pylab.imshow(IAB1, origin='lower')
    pylab.title('Mutual Information')

#===============================================================================
# Test
#===============================================================================
def Test():
    """
    """
#    path1='C:/Users/bryan/bryan-code/trunk/Images/Test_vert.png'
#    path2='C:/Users/bryan/bryan-code/trunk/Images/Test_horiz.png'
    path1 = 'C:/Users/bryan/bryan-code/trunk/Images/ucharim080-LAT.png'
    path2 = 'C:/Users/bryan/bryan-code/trunk/Images/CalibratedDRRImage.png'
    img1 = Image.open(path1, 'r')
    img2 = Image.open(path2, 'r')
    roi = [[210, 187], [277, 370]]
    #cost=GetValue(img1,img2,roi,'gd')
    cost = GetValue(img1, img2, method='ncc', gsigma=1.4, roi=roi, ret_histo=False, bins=256, debug=True)
    print "Metric Done!"
    return cost

if __name__ == "__main__":
    Test()
    